{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "159d321c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "from openai.types.responses import ResponseTextDeltaEvent\n",
    "from pydantic import BaseModel, Field\n",
    "from agents import Agent, trace, Runner, WebSearchTool\n",
    "from dotenv import load_dotenv\n",
    "from IPython.display import display, Markdown\n",
    "import asyncio\n",
    "from agents.model_settings import ModelSettings\n",
    "\n",
    "load_dotenv(override=True)\n",
    "\n",
    "class Evaluation(BaseModel):\n",
    "    is_satisfactory: bool = Field(description=\"Whether the user's answers are satisfactory to our questions.\")\n",
    "    feedback: str = Field(description=\"Feedback on the user's answers.\")\n",
    "    score: int = Field(description=\"Score between 0 and 1 based on the quality of the answers.\")\n",
    "\n",
    "class WebSearchItem(BaseModel):\n",
    "    reason: str = Field(description=\"Your reasoning for why this search is important to the query.\")\n",
    "\n",
    "    query: str = Field(description=\"The search term to use for the web search.\")\n",
    "\n",
    "class WebSearchPlan(BaseModel):\n",
    "    searches: list[WebSearchItem] = Field(description=\"A list of web searches to perform to best answer the query.\")\n",
    "\n",
    "class ReportData(BaseModel):\n",
    "    short_summary: str = Field(description=\"A short 2-3 sentence summary of the findings.\")\n",
    "\n",
    "    markdown_report: str = Field(description=\"The final report\")\n",
    "\n",
    "    follow_up_questions: list[str] = Field(description=\"Suggested topics to research further\")\n",
    "\n",
    "HOW_MANY_SEARCHES = 3\n",
    "\n",
    "planner_instructions= f\"You are a helpful research assistant. Given a query, come up with a set of web searches \\\n",
    "to perform to best answer the query. Output {HOW_MANY_SEARCHES} terms to query for.\"\n",
    "planner_agent = Agent(\n",
    "    name=\"PlannerAgent\",\n",
    "    instructions=planner_instructions,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=WebSearchPlan,\n",
    ")\n",
    "\n",
    "writer_instructions = (\n",
    "    \"You are a senior researcher tasked with writing a cohesive report for a research query. \"\n",
    "    \"You will be provided with the original query, and some initial research done by a research assistant.\\n\"\n",
    "    \"You should first come up with an outline for the report that describes the structure and \"\n",
    "    \"flow of the report. Then, generate the report and return that as your final output.\\n\"\n",
    "    \"The final output should be in markdown format, and it should be lengthy and detailed. Aim \"\n",
    "    \"for 5-10 pages of content, at least 1000 words.\"\n",
    ")\n",
    "writer_agent = Agent(\n",
    "    name=\"WriterAgent\",\n",
    "    instructions=writer_instructions,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=ReportData,\n",
    ")\n",
    "\n",
    "search_instructions = \"You are a research assistant. Given a search term, you search the web for that term and \\\n",
    "produce a concise summary of the results. The summary must 2-3 paragraphs and less than 300 \\\n",
    "words. Capture the main points. Write succintly, no need to have complete sentences or good \\\n",
    "grammar. This will be consumed by someone synthesizing a report, so it's vital you capture the \\\n",
    "essence and ignore any fluff. Do not include any additional commentary other than the summary itself.\"\n",
    "search_agent = Agent(\n",
    "    name=\"Search agent\",\n",
    "    instructions=search_instructions,\n",
    "    tools=[WebSearchTool(search_context_size=\"low\")],\n",
    "    model=\"gpt-4o-mini\",\n",
    "    model_settings=ModelSettings(tool_choice=\"required\"),\n",
    ")\n",
    "\n",
    "eval_instructions = \"\"\"\n",
    "You need to evaluate whether the user's answers are satisfactory to our questions. \n",
    "If not, you need to tell the user where it is lacking.\n",
    "Do not be strict, if the answers are enough to continue with the research, you can say it is satisfactory.\n",
    "\"\"\"\n",
    "evaluator_agent = Agent(\n",
    "    name=\"Evaluator\",\n",
    "    instructions=eval_instructions,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=Evaluation,\n",
    ")\n",
    "clarifying_agent = Agent(\n",
    "    name=\"Clarifying Questions Agent\",\n",
    "    instructions=\"You are a helpful assistant that asks three clarifying questions to the user.\",\n",
    "    model=\"gpt-4o-mini\",\n",
    ")\n",
    "\n",
    "async def plan_searches(query: str):\n",
    "    \"\"\" Use the planner_agent to plan which searches to run for the query \"\"\"\n",
    "    print(\"Planning searches...\")\n",
    "    result = await Runner.run(planner_agent, f\"Query: {query}\")\n",
    "    print(f\"Will perform {len(result.final_output.searches)} searches\")\n",
    "    return result.final_output\n",
    "\n",
    "async def search(item: WebSearchItem):\n",
    "    \"\"\" Use the search agent to run a web search for each item in the search plan \"\"\"\n",
    "    input = f\"Search term: {item.query}\\nReason for searching: {item.reason}\"\n",
    "    result = await Runner.run(search_agent, input)\n",
    "    return result.final_output\n",
    "\n",
    "async def perform_searches(search_plan: WebSearchPlan):\n",
    "    \"\"\" Call search() for each item in the search plan \"\"\"\n",
    "    print(\"Searching...\")\n",
    "    tasks = [asyncio.create_task(search(item)) for item in search_plan.searches]\n",
    "    results = await asyncio.gather(*tasks)\n",
    "    print(\"Finished searching\")\n",
    "    return results\n",
    "\n",
    "async def write_report(query: str, search_results: list[str]):\n",
    "    \"\"\" Use the writer agent to write a report based on the search results\"\"\"\n",
    "    print(\"Thinking about report...\")\n",
    "    input = f\"Original query: {query}\\nSummarized search results: {search_results}\"\n",
    "    result = await Runner.run(writer_agent, input)\n",
    "    print(\"Finished writing report\")\n",
    "    return result.final_output\n",
    "\n",
    "async def run_research(query: str, questions: str, answers: str):\n",
    "    \"\"\" Run the full research pipeline: plan searches, perform searches, write report, and send email \"\"\"\n",
    "    with trace(\"Research trace\"):\n",
    "        print(\"Starting research...\")\n",
    "        input = f\"Query: {query}\\nQuestions: {questions}\\nAnswers: {answers}\"\n",
    "        search_plan = await plan_searches(input)\n",
    "        search_results = await perform_searches(search_plan)\n",
    "        report = await write_report(query, search_results)\n",
    "        print(\"Hooray!\")\n",
    "        return report.markdown_report\n",
    "\n",
    "async def call_questions_stream(topic: str):\n",
    "    \"\"\"Streaming version - THIS IS THE FIX\"\"\"\n",
    "    yield gr.update(visible=True, value=\"\")\n",
    "    with trace(\"Clarifying Questions\"):\n",
    "        full_text = \"\"\n",
    "        result = Runner.run_streamed(clarifying_agent, input=topic)\n",
    "        async for event in result.stream_events():\n",
    "            if event.type == \"raw_response_event\" and isinstance(event.data, ResponseTextDeltaEvent):\n",
    "                full_text+=event.data.delta\n",
    "                yield full_text\n",
    "\n",
    "async def evaluate_and_research(topic: str, questions: str, answers: str):\n",
    "    \"\"\"Evaluate answers first, then run research if satisfactory\"\"\"\n",
    "    # Run evaluator\n",
    "    input = f\"Topic:\\n{topic}\\n\\nQuestions:\\n{questions}\\n\\nAnswers:\\n{answers}\\n\\n\\\n",
    "        Criteria: Evaluate if the answers provide enough focus and specificity for conducting targeted research.\"\n",
    "    evaluation = await Runner.run(evaluator_agent, input=input)\n",
    "    evaluation_result = evaluation.final_output\n",
    "    if not evaluation_result.is_satisfactory and evaluation_result.score<0.5:\n",
    "        # Return feedback if not satisfactory\n",
    "        feedback_message = f\"❌ **Answers need improvement:**\\n\\n{evaluation_result.feedback}\\n\\nPlease revise your answers and click the button again.\"\n",
    "        yield gr.update(visible=True, value=feedback_message)\n",
    "        return  # Stop here if not satisfactory\n",
    "    \n",
    "    # If satisfactory, proceed with research\n",
    "    success_message = \"✅ **Answers are satisfactory! Starting research...**\\n\\n\"\n",
    "    yield gr.update(visible=True, value=success_message)\n",
    "    \n",
    "    # Run the actual research\n",
    "    research_result = await run_research(topic, questions, answers)\n",
    "    yield gr.update(visible=True, value=research_result)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dae3d95",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gr.Blocks() as demo:\n",
    "    gr.Markdown(\"# 🔬 Deep Research Assistant\")\n",
    "    gr.Markdown(\"Enter your research topic below. I'll ask clarifying questions to better understand your needs, then conduct comprehensive research.\")\n",
    "    \n",
    "    with gr.Row():\n",
    "        with gr.Column(scale=1):\n",
    "            topic = gr.Textbox(\n",
    "                label=\"Research Topic\",\n",
    "                placeholder=\"e.g., What is the best AI Agent framework in 2025?\",\n",
    "                lines=2\n",
    "            )\n",
    "            clarify_btn = gr.Button(\"🤔 Get Clarifying Questions\", variant=\"primary\", size=\"lg\")\n",
    "            \n",
    "            questions = gr.TextArea(visible=False, label=\"Clarifying Questions\")\n",
    "            answers = gr.Textbox(\n",
    "                visible=False,\n",
    "                label=\"Your Answers\",\n",
    "                placeholder=\"Please answer the questions above to help refine the research scope...\",\n",
    "                lines=5\n",
    "            )\n",
    "            \n",
    "            research_btn = gr.Button(\"🚀 Start Deep Research\", visible=False, variant=\"primary\", size=\"lg\")\n",
    "        \n",
    "        with gr.Column(scale=1):\n",
    "            report = gr.Markdown(visible=False, label=\"Research Report\", header_links=True,\n",
    "                    show_copy_button=True,line_breaks=True,sanitize_html=True, height=500,)\n",
    "\n",
    "    clarify_btn.click(\n",
    "        fn=call_questions_stream,\n",
    "        inputs=topic,\n",
    "        outputs=[questions]\n",
    "    ).then(\n",
    "        fn=lambda: (gr.update(visible=True, interactive=True), gr.update(visible=True)),\n",
    "        inputs=None,\n",
    "        outputs=[answers, research_btn]\n",
    "    )\n",
    "\n",
    "    \n",
    "    research_btn.click(\n",
    "        fn=evaluate_and_research,\n",
    "        inputs=[topic, questions, answers],\n",
    "        outputs=[report]\n",
    "    )\n",
    "\n",
    "\n",
    "demo.launch(inbrowser=True, inline=False)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
